import os
import cv2
import easyocr
from tqdm import tqdm
import pandas as pd
import numpy as np

# 初始化EasyOCR阅读器
reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)  # 确保关闭GPU

def read_image_safe(path):
    """安全读取图片（支持中文路径）"""
    try:
        with open(path, 'rb') as f:
            img_bytes = f.read()
        img_buffer = np.frombuffer(img_bytes, dtype=np.uint8)
        img = cv2.imdecode(img_buffer, cv2.IMREAD_COLOR)
        return img
    except:
        return None

def batch_ocr(image_folder, output_file='output.csv'):
    results = []
    image_list = [
        os.path.join(image_folder, f) 
        for f in os.listdir(image_folder)
        if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))
    ]
    
    for image_path in tqdm(image_list, desc="Processing"):
        try:
            # 安全读取图片
            img = read_image_safe(image_path)
            if img is None:
                print(f"图片读取失败：{image_path}")
                continue
            
            # 转换为RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 执行OCR
            result = reader.readtext(img, detail=0)
            full_text = ' '.join(result)
            
            results.append({
                'image_path': image_path,
                'text': full_text
            })
            print(full_text)
        except Exception as e:
            print(f"Error processing {image_path}: {str(e)}")
            continue
    
    # 保存结果
    pd.DataFrame(results).to_csv(output_file, index=False, encoding='utf-8-sig')
    print(f"结果已保存至 {output_file}")

# 执行批量识别
batch_ocr(image_folder=r'D:/Desk/0312-炳哥尺寸转换/尺寸标注/标注信息')